
import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from skimage.feature import graycomatrix, graycoprops
from skimage.morphology import skeletonize
from skimage import exposure
from skimage.util.dtype import img_as_bool
from StructureTensor import cal_structure_tensor


def get_files_by_ext(directory, ext):
    return [f for f in os.listdir(directory) if f.endswith(ext)]

def remove_shadow(img):
    # 基于亮度的阴影检测与去除
    rgb = cv2.cvtColor(cv2.cvtColor(img, cv2.COLOR_GRAY2BGR), cv2.COLOR_BGR2LAB)
    l_channel = rgb[:,:,0]
    _, shadow_mask = cv2.threshold(l_channel, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    shadow_mask = cv2.bitwise_not(shadow_mask)
    
    # 阴影区域亮度补偿
    non_shadow = cv2.bitwise_and(img, img, mask=cv2.bitwise_not(shadow_mask))
    shadow_area = cv2.bitwise_and(img, img, mask=shadow_mask)
    compensated_shadow_area = cv2.addWeighted(shadow_area, 1.5, np.zeros_like(shadow_area), 0, 30)
    return cv2.bitwise_or(non_shadow, compensated_shadow_area)


def remove_illumination_trend(img):
    # 先去除阴影
    img = remove_shadow(img)
    
    # CLAHE增强对比度
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
    enhanced = clahe.apply(img)
    
    # 形态学Top-Hat变换去除低频亮度趋势
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (25,25))
    background = cv2.morphologyEx(enhanced, cv2.MORPH_OPEN, kernel)
    return cv2.subtract(enhanced, background)


def postprocess_mask(mask):
    # 形态学闭运算填充空洞
    kernel = np.ones((5,5), np.uint8)
    closed = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    
    # 连通域分析去除小噪声
    _, labels = cv2.connectedComponents(closed)
    unique, counts = np.unique(labels, return_counts=True)
    for label, count in zip(unique[1:], counts[1:]):  # 跳过背景
        if count < 100:  # 面积阈值
            closed[labels == label] = 0
    return closed

# ---------- 1. 读取与预处理 ----------
# image_file = "7-17b_rotated_columns_1.jpg"
def image_filter(image_file):
    img = cv2.imread(image_file, 0)           # 灰度图
    # s0 = cal_structure_tensor(img, sigma=2, rho=2, plt_min=0, plt_max=255, model='2d')

    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    enhanced = clahe.apply(img)
    # Top-Hat变换去除亮度趋势
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (21,21))
    background = cv2.morphologyEx(enhanced, cv2.MORPH_OPEN, kernel)
    # 添加中值滤波压制随机噪声
    trend_removed = enhanced - background
    trend_removed = cv2.medianBlur(trend_removed, ksize=3)
    

    # 自适应阈值 → 二值纹理掩膜
    _, bw = cv2.threshold(trend_removed.astype(np.uint8), 0, 255,
                        cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    # 方向纹理滤波：压制垂直，增强水平
     # 2. 水平纹理增强
    horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 1))  # 水平结构元素
    bw = cv2.morphologyEx(bw, cv2.MORPH_CLOSE, horizontal_kernel)  # 增强水平纹理
    
    # # 1. 垂直纹理压制
    # vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 5))  # 垂直结构元素
    # bw = cv2.morphologyEx(bw, cv2.MORPH_OPEN, vertical_kernel)  # 去除垂直纹理
    
   

    cv2.imwrite(image_file.replace('.jpg', '_filter.png'), bw)
    print(image_file.replace('.jpg', '_filter.png'))

    # ---------- 3. 纹理基元提取（裂隙/层理骨架） ----------
    skeleton = skeletonize(bw//255).astype(np.uint8) * 255

    # ---------- 4. 计算结构参数 ----------
    # 4-a) 主要方向（Hough）
    # 优化霍夫变换参数以更好检测长条形结构
    lines = cv2.HoughLinesP(skeleton, 1, np.pi/180,
                            threshold=30, minLineLength=50, maxLineGap=10)
    angles = []
    if lines is not None:
        for x1,y1,x2,y2 in lines[:,0]:
            angles.append(np.arctan2(y2-y1, x2-x1) * 180/np.pi)
    mean_angle = np.mean(angles) if angles else 0

    # 4-b) 纹理对比度（GLCM）
    glcm = graycomatrix(img, distances=[5], angles=[0, np.pi/4, np.pi/2, 3*np.pi/4],
                        levels=256, symmetric=True, normed=True)
    contrast = graycoprops(glcm, 'contrast').mean()

    #---------- 5. 结果可视化 ----------
    plt.figure(figsize=(12,4))
    plt.title(image_file)
    plt.subplot(131); plt.imshow(img, cmap='gray');   plt.title('Original')
    plt.subplot(132); plt.imshow(bw, cmap='gray');    plt.title('Enhanced texture')
    plt.subplot(133); plt.imshow(skeleton, cmap='gray'); plt.title('Skeleton / Layering')
    plt.suptitle(f'Mean orientation = {mean_angle:.1f}°, GLCM contrast = {contrast:.2f}')
    plt.tight_layout(); plt.show()

if __name__ == "__main__":
    files = get_files_by_ext('./texture_template',"jpg")
    for file_image in files:
        image_filter('./texture_template/'+file_image)
        